import cv2
import torch
import os
import uuid

def process_image(image_path):
    # 创建 outputs 文件夹（如果不存在）
    output_dir = "./output"
    os.makedirs(output_dir, exist_ok=True)
    
    # 生成唯一的输出文件名
    output_filename = f"processed_image_{uuid.uuid4().hex}.jpg"
    output_path = os.path.join(output_dir, output_filename)

    model = torch.hub.load('./yolov5', 'custom', path='./models/hats_best.pt', source='local')

    # 读取图片
    frame = cv2.imread(image_path)

    # 进行目标检测
    results = model(frame)

    # 渲染结果
    rendered_frame = results.render()[0]

    # 保存结果
    cv2.imwrite(output_path, rendered_frame)

    # 显示结果
    cv2.imshow('YOLOv5 Detection', rendered_frame)
    cv2.waitKey(0)
    cv2.destroyAllWindows()